import librosa  # 导入librosa库用于音频处理
import numpy as np
from funasr import AutoModel
import os
import jieba

# 新增获取当前文件所在目录路径
current_dir = os.path.dirname(os.path.abspath(__file__))
# 设置jieba字典路径（将dict.txt放在models/iic目录下）
# jieba.set_dictionary(os.path.join(current_dir, "models/iic/jieba_dict.txt"))

# 如果已有缓存文件需要指定缓存位置（可选）
jieba.dt.cache_file = os.path.join(current_dir, "models/iic/jieba.cache")
# 修改为使用动态路径组合
model_dir = os.path.join(current_dir, "models/iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")
vad_model_dir = os.path.join(current_dir, "models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch")
punc_model_dir = os.path.join(current_dir, "models/iic/punc_ct-transformer_cn-en-common-vocab471067-large")
model = AutoModel(model=model_dir, vad_model=vad_model_dir, vad_kwargs={"max_single_segment_time": 30000}, device="cuda:0", disable_update=True)
punc_model = AutoModel(model=punc_model_dir, disable_update=True, device="cuda:0")


def splice_chunks(audio_cache, audio_chunk):
    if audio_chunk is not None:
        sr, y = audio_chunk
        # ... 保留预处理逻辑 ...
        if y.ndim > 1:
            y = y.mean(axis=1)
        y = y.astype(np.float32)
        # 添加空数组检查
        if len(y) == 0:
            return audio_cache  # 返回上一次的结果
        # 添加非零检查
        if np.max(np.abs(y)) == 0:
            return audio_cache  # 返回上一次的结果
        y /= np.max(np.abs(y))
        if y is None:
            return audio_cache
        speech = y  # 提取音频数据
        # 持续累积音频数据
        if audio_cache is not None:
            audio_cache = np.concatenate([audio_cache, speech])
        else:
            audio_cache = speech
        return audio_cache


def asr_decode(audio_cache):
    if audio_cache is not None and len(audio_cache) > 0:
        speech = librosa.resample(y=audio_cache, orig_sr=48000, target_sr=16000)
        result = model.generate(speech, use_itn=True)
        text_output = result[0]['text'] if result else ""
        text_output = punc_model.generate(text_output)[0]['text']
        return text_output + ""  # 清空state并返回结果
    return ""
